import torch
import torch.nn as nn


class SimpleLLM(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_size):
        super(SimpleLLM, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_size, batch_first=True)
        self.fc = nn.Linear(hidden_size, vocab_size)

    def forward(self, x):
        x = self.embedding(x)
        x, _ = self.lstm(x)
        if x.dim() == 2:
            x = self.fc(x)
        elif x.dim() == 3:
            x = self.fc(x[:, -1, :])
        return x